package com.swak.zookeeper;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentMap;

import com.swak.Constants;
import com.swak.codec.Encodes;
import com.swak.registry.ChildListener;
import com.swak.registry.FailbackRegistry;
import com.swak.registry.NotifyListener;
import com.swak.registry.State;
import com.swak.registry.URL;
import com.swak.utils.CollectionUtils;
import com.swak.utils.Sets;

/**
 * 如果是订阅：不关注中间的目录，必须对应到具体的子节点
 * 默认是支持三层；现在改进为支持四层；
 * 
 * @author lifeng
 */
public class ZookeeperRegistry extends FailbackRegistry {

	private ZookeeperService zookeeperService;
	private final ConcurrentMap<URL, Set<String>> anyModules = new ConcurrentHashMap<>();
	private final ConcurrentMap<URL, Set<String>> anyServices = new ConcurrentHashMap<>();
	private final ConcurrentMap<URL, ConcurrentMap<NotifyListener, ChildListener>> zkListeners = new ConcurrentHashMap<>();

	public ZookeeperRegistry(URL url, ZookeeperService zookeeperService) {
		super(url);
		this.zookeeperService = zookeeperService;
		this.zookeeperService.addStateListener(state -> {
			if (state == State.RECONNECTED) {
				logger.warn(
						"Trying to fetch the latest urls, in case there're provider changes during connection loss.\n"
								+ " Since ephemeral ZNode will not get deleted for a connection lose, "
								+ "there's no need to re-register url of this instance.");
				this.fetchLatestAddresses();
			} else if (state == State.NEW_SESSION_CREATED) {
				logger.warn("Trying to re-register urls and re-subscribe listeners of this instance to registry...");
				try {
					ZookeeperRegistry.this.recover();
				} catch (Exception e) {
					logger.error(e.getMessage(), e);
				}
			} else if (state == State.SESSION_LOST) {
				logger.warn("Url of this instance will be deleted from registry soon. "
						+ "Dubbo client will try to re-register once a new session is created.");
			} else if (state == State.SUSPENDED) {

			} else if (state == State.CONNECTED) {

			}
		});
	}

	@Override
	public void doRegister(URL url) {
		this.create(url.toUrlPath(), true);
	}

	@Override
	public void doUnregister(URL url) {
		this.delete(url.toUrlPath());
	}

	protected void create(String path, boolean ephemeral) {
		zookeeperService.create(path, ephemeral);
	}

	protected void delete(String path) {
		zookeeperService.delete(path);
	}

	@Override
	public void doSubscribe(URL url, NotifyListener listener) {
		if (Constants.Subscribe_Any_Value.equals(url.getModule())) {
			this.doAnyModuleSubscribe(url, listener);
		} else if (Constants.Subscribe_Any_Value.equals(url.getService())) {
			this.doAnyServiceSubscribe(url, listener);
		} else {
			this.doCategorySubscribe(url, listener);
		}
	}

	private void doAnyModuleSubscribe(URL url, NotifyListener listener) {
		ConcurrentMap<NotifyListener, ChildListener> listeners = zkListeners.computeIfAbsent(url,
				k -> new ConcurrentHashMap<>());
		ChildListener zkListener = listeners.computeIfAbsent(listener, k -> (parentPath, currentChilds) -> {
			for (String child : currentChilds) {
				child = Encodes.urlDecode(child);
				Set<String> services = anyModules.get(url);
				if (services == null) {
					services = anyModules.computeIfAbsent(url, (s) -> {
						return Sets.newConcurrentSet();
					});
				}
				if (!services.contains(child)) {
					services.add(child);
					subscribe(URL.valueOf(url).addParameter(Constants.Module_Key, child), k);
				}
			}
		});
		this.create(url.toRootPath(), false);
		List<String> currentChilds = zookeeperService.addChildListener(url.toRootPath(), zkListener);
		if (CollectionUtils.isNotEmpty(currentChilds)) {
			for (String child : currentChilds) {
				child = Encodes.urlDecode(child);
				Set<String> services = anyModules.get(url);
				if (services == null) {
					services = anyModules.computeIfAbsent(url, (s) -> {
						return Sets.newConcurrentSet();
					});
				}
				if (!services.contains(child)) {
					services.add(child);
					subscribe(URL.valueOf(url).addParameter(Constants.Module_Key, child), listener);
				}
			}
		}
	}

	private void doAnyServiceSubscribe(URL url, NotifyListener listener) {
		ConcurrentMap<NotifyListener, ChildListener> listeners = zkListeners.computeIfAbsent(url,
				k -> new ConcurrentHashMap<>());
		ChildListener zkListener = listeners.computeIfAbsent(listener, k -> (parentPath, currentChilds) -> {
			for (String child : currentChilds) {
				child = Encodes.urlDecode(child);
				Set<String> services = anyServices.get(url);
				if (services == null) {
					services = anyServices.computeIfAbsent(url, (s) -> {
						return Sets.newConcurrentSet();
					});
				}
				if (!services.contains(child)) {
					services.add(child);
					subscribe(URL.valueOf(url).addParameter(Constants.Service_Key, child), k);
				}
			}
		});
		this.create(url.toModulePath(), false);
		List<String> currentChilds = zookeeperService.addChildListener(url.toModulePath(), zkListener);
		if (CollectionUtils.isNotEmpty(currentChilds)) {
			for (String child : currentChilds) {
				child = Encodes.urlDecode(child);
				Set<String> services = anyServices.get(url);
				if (services == null) {
					services = anyServices.computeIfAbsent(url, (s) -> {
						return Sets.newConcurrentSet();
					});
				}
				if (!services.contains(child)) {
					services.add(child);
					subscribe(URL.valueOf(url).addParameter(Constants.Service_Key, child), listener);
				}
			}
		}
	}

	private void doCategorySubscribe(URL url, NotifyListener listener) {
		List<URL> urls = new ArrayList<>();
		for (String path : url.toCategoryPaths()) {
			ConcurrentMap<NotifyListener, ChildListener> listeners = zkListeners.computeIfAbsent(url,
					k -> new ConcurrentHashMap<>());
			ChildListener zkListener = listeners.computeIfAbsent(listener, k -> (parentPath,
					currentChilds) -> ZookeeperRegistry.this.notify(url, k, toUrls(url, currentChilds)));
			this.create(path, false);
			List<String> children = zookeeperService.addChildListener(path, zkListener);
			if (children != null) {
				urls.addAll(toUrls(url, children));
			}
		}
		this.notify(url, listener, urls);
	}

	private List<URL> toUrls(URL consumer, List<String> providers) {
		List<URL> urls = new ArrayList<>();
		if (CollectionUtils.isNotEmpty(providers)) {
			for (String provider : providers) {
				URL url = URL.valueOf(Encodes.urlDecode(provider));
				if (consumer.match(url)) {
					urls.add(url);
				}
			}
		}
		return urls;
	}

	@Override
	public void doUnsubscribe(URL url, NotifyListener listener) {
		ConcurrentMap<NotifyListener, ChildListener> listeners = zkListeners.get(url);
		if (listeners != null) {
			ChildListener zkListener = listeners.get(listener);
			if (zkListener != null) {
				if (Constants.Subscribe_Any_Value.equals(url.getService())) {
					zookeeperService.removeChildListener(url.toRootPath(), zkListener);
				} else {
					for (String path : url.toCategoryPaths()) {
						zookeeperService.removeChildListener(path, zkListener);
					}
				}
			}
		}
	}

	/**
	 * When zookeeper connection recovered from a connection loss, it need to fetch
	 * the latest provider list. re-register watcher is only a side effect and is
	 * not mandate.
	 */
	private void fetchLatestAddresses() {
		// subscribe
		Map<URL, Set<NotifyListener>> recoverSubscribed = new HashMap<URL, Set<NotifyListener>>(getSubscribed());
		if (!recoverSubscribed.isEmpty()) {
			if (logger.isInfoEnabled()) {
				logger.info("Fetching the latest urls of " + recoverSubscribed.keySet());
			}
			for (Map.Entry<URL, Set<NotifyListener>> entry : recoverSubscribed.entrySet()) {
				URL url = entry.getKey();
				for (NotifyListener listener : entry.getValue()) {
					addFailedSubscribed(url, listener);
				}
			}
		}
	}
}
